# json_merge.py
# usage:
# python ./src/tools/json_merge.py --src1 "./result/results.part1.jsonl" --src2 "./result/results.part2.jsonl" --des "./result/results.jsonl"

import argparse
import json
import sys

def read_jsonl(file_path):
    """读取 JSONL 文件，返回以 id 为键的字典"""
    data = {}
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            for line_num, line in enumerate(f, 1):
                line = line.strip()
                if not line:
                    continue  # 跳过空行
                try:
                    obj = json.loads(line)
                    if 'id' not in obj:
                        print(f"警告: 第 {line_num} 行缺少 'id' 字段，跳过: {line}")
                        continue
                    data[obj['id']] = obj
                except json.JSONDecodeError as e:
                    print(f"错误: 第 {line_num} 行 JSON 解析失败: {e}")
                    sys.exit(1)
    except FileNotFoundError:
        print(f"错误: 文件未找到: {file_path}")
        sys.exit(1)
    except Exception as e:
        print(f"错误: 读取文件时发生异常: {e}")
        sys.exit(1)
    return data

def write_jsonl(data, file_path):
    """将数据按 id 排序后写入 JSONL 文件"""
    try:
        with open(file_path, 'w', encoding='utf-8') as f:
            # 按 id 排序
            sorted_items = sorted(data.items(), key=lambda x: x[0])
            for _id, obj in sorted_items:
                f.write(json.dumps(obj, ensure_ascii=False) + '\n')
    except Exception as e:
        print(f"错误: 写入文件失败: {e}")
        sys.exit(1)

def main():
    parser = argparse.ArgumentParser(description="合并两个 JSONL 文件，按 id 排序，src2 优先")
    parser.add_argument('--src1', required=True, help='第一个源文件路径')
    parser.add_argument('--src2', required=True, help='第二个源文件路径')
    parser.add_argument('--des', required=True, help='输出目标文件路径')

    args = parser.parse_args()

    # 读取文件
    src1_data = read_jsonl(args.src1)
    src2_data = read_jsonl(args.src2)

    # 合并：src2 优先
    merged = src1_data.copy()
    merged.update(src2_data)  # src2 会覆盖 src1 中相同 id 的记录

    # 写入目标文件
    write_jsonl(merged, args.des)
    print(f"成功合并 {len(src1_data)} + {len(src2_data)} 条记录，输出 {len(merged)} 条记录到 {args.des}")

if __name__ == '__main__':
    main()